import { create } from "zustand";
import { persist } from "zustand/middleware";
import { immer } from "zustand/middleware/immer";
import {
    ChatResponse,
    MessageRole,
    ChatInterface,
    Message,
    ConversationInterface,
} from "../../core/types/chat";
import { sendMessage } from "../../core/services/ollama_service";
import { Conversations } from "../../core/types/chat";
import { createTimeStamp } from "../../util/crypto";
import { useModal } from "./useModal";
import { useUpload } from "./useUpload";
import { zustandPersistStorage } from "../../util/storage";
import { mcpTools } from "../../core/services/mcp_service";
import { jsonClone, prettyJson } from "../../util/tools";
import { ConfigService } from "../../core/services/config_service";
import { WritableDraft } from "immer";

const USER_ROLE: MessageRole = "user";
const ASSIST_ROLE: MessageRole = "assistant";
const TOOL_ROLE: MessageRole = "tool";
const SYSTEM_ROLE: MessageRole = "system";

const SYSTEM_PROMPT = `你是一个AI智能对话助手，你遵守以下规则：
- 默认你的用户是中国大陆居民
- 默认你的用户年龄区间是18-65岁
- 用户的个人隐私信息无论任何情况下都不能泄露
- 尽量使用中文回答
- 尽量使用简洁的语言回答用户的问题
- 有工具可以完成用户命令的，则你必须使用工具，不需要用户授权`;

// please update this value when you add new initial messages
const DEFAULT_MESSAGES_LENGTH = 2;
function createNewConversation() {
    const systemMessage: Message = {
        id: createTimeStamp(),
        role: SYSTEM_ROLE,
        timestamp: Date.now(),
        content: SYSTEM_PROMPT,
    };
    const assistantMessage: Message = {
        id: createTimeStamp(),
        role: ASSIST_ROLE,
        timestamp: Date.now(),
        content: "你好，我是你的AI智能对话助手，有什么可以帮到你吗？",
    };
    return {
        id: createTimeStamp(),
        title: "新对话",
        createdAt: Date.now(),
        updatedAt: Date.now(),
        // update DEFAULT_MESSAGES_LENGTH when you add new initial messages
        messages: [systemMessage, assistantMessage],
    };
}

const chatStore = immer<ChatInterface>((set, get) => {
    return {
        currentConversationId: "",
        setCurrentConversationId: (conversationId: string) => {
            set((state) => {
                state.currentConversationId = conversationId;
            });
        },

        conversations: [],
        setConversations: (conversations: Conversations) => {
            set((state) => {
                state.conversations = conversations;
            });
        },
        getCurrentConversation() {
            const state = get();
            if (!state) {
                console.warn("no state");
                return;
            }
            const conversations = state.conversations;
            const currentConversationId = state.currentConversationId;
            return conversations.find(
                (conversation) => conversation.id === currentConversationId,
            );
        },
        isHydrated: false,
        initialize: () => {
            set((stateDraft) => {
                if (!stateDraft.isHydrated) {
                    return;
                }
                if (stateDraft.conversations.length > 0) {
                    return;
                }
                stateDraft.isLoading = false;
                const newConversation = createNewConversation();
                stateDraft.conversations = [newConversation];
                stateDraft.currentConversationId = newConversation.id;
            });
        },
        isLoading: false,
        setIsLoading: (isLoading: boolean) => {
            set((state) => {
                state.isLoading = isLoading;
            });
        },
        clearConversations: () => {
            useUpload.getState().clearFiles();
            set((state) => {
                state.conversations = [];
                state.currentConversationId = "";
                state.initialize();
            });
        },
        currentInput: "",
        setCurrentInput: (input: string) => {
            set((state) => {
                state.currentInput = input;
            });
        },
        sendMessage: async (content: string, role: MessageRole = "user") => {
            if (get().isLoading) {
                return;
            }
            const i = get().conversations.findIndex(
                (e) => e.id === get().currentConversationId,
            );
            if (i === -1) {
                console.warn("no current conversation");
                return;
            }
            const configService = ConfigService.getInstance();
            const config = configService.getConfig();
            const modelName = config.selectedModel;
            if (!modelName) {
                useModal
                    .getState()
                    .alert("请先选择模型，在右上角点击“配置”图标");
                return;
            }

            set((state) => {
                state.isLoading = true;
                const messages = state.conversations[i].messages;
                state.currentInput = "";
                let newMessage: Message | undefined;
                newMessage = {
                    id: createTimeStamp(),
                    role,
                    content,
                    timestamp: Date.now(),
                    images: [],
                };
                if (!newMessage) {
                    throw new Error("no new message");
                }
                const { files, markFileUsed } = useUpload.getState();
                if (Array.isArray(files)) {
                    files.forEach((fileInfo) => {
                        if (fileInfo.messageId) {
                            return;
                        }
                        if (newMessage.images) {
                            newMessage.images.push(
                                fileInfo.base64.split(",")[1],
                            );
                        }
                        markFileUsed(fileInfo.id, newMessage.id);
                    });
                }
                messages.push(newMessage);
            });
            let messagesArg = jsonClone(get().conversations[i].messages);
            const arg: { messages: []; tools: [] } = {
                messages: messagesArg,
                tools: [],
            };
            const modelInfo = await configService.getModelInfo(modelName);
            if (modelInfo?.capabilities?.includes("tools")) {
                const tools = await mcpTools;
                if (!tools) {
                    useModal.getState().alert("获取工具列表失败");
                    return;
                }
                arg.tools = tools;
            }

            sendMessage((chunk: ChatResponse) => {
                onStream(chunk, set);
            }, arg).catch((error) => {
                useModal.getState().alert(error.message);
            });
        },
        newConversation: () => {
            set((state) => {
                const currentConversation = state.getCurrentConversation();
                if (
                    currentConversation?.messages.length ===
                    DEFAULT_MESSAGES_LENGTH
                ) {
                    // current conversation is empty, no need to create a new one
                    return;
                }
                // find a empty conversation, if not, create a new one
                const emptyConversation = state.conversations.find(
                    (conversation) => !conversation.messages.length,
                );
                if (emptyConversation) {
                    state.currentConversationId = emptyConversation.id;
                    return;
                }
                // create a new conversation
                const newConversation = createNewConversation();
                const conversations = state.conversations.slice(0, 9);
                conversations.unshift(newConversation);
                state.conversations = conversations;
                state.currentConversationId = newConversation.id;
            });
        },
        deleteConversation: (conversationId: string) => {
            set((state) => {
                state.conversations = state.conversations.filter(
                    (conversation) => conversation.id !== conversationId,
                );
            });
            useUpload.getState().removeConversationFiles(conversationId);
        },
    };
});

function onStream(
    chunk: ChatResponse,
    set: (
        nextStateOrUpdater: (state: WritableDraft<ChatInterface>) => void,
    ) => void,
) {
    set((state: WritableDraft<ChatInterface>) => {
        const currentConversation = state.conversations.find(
            (conversation: ConversationInterface) =>
                conversation.id === state.currentConversationId,
        );
        if (!currentConversation) {
            throw new Error("no current conversation");
        }
        const messages = currentConversation.messages;
        // get last one message
        let lastMessage = messages[messages.length - 1];
        const chunkMessage = chunk.message;
        if (!chunkMessage) {
            throw new Error("no chunk message");
        }
        const chunkContent = chunkMessage.content || "";
        if (chunkMessage.tool_calls) {
            // 当AI调用工具时，不用向用户输出内容，只输出工具调用的方法和参数
            messages.push({
                id: createTimeStamp(),
                role: SYSTEM_ROLE,
                content:
                    `以下是你即将调用工具的方法和参数，不要思考，直接输出用户能理解的调用说明，` +
                    `例如“正在调用某某工具，参数为某某某，或者无参数”：\n` +
                    prettyJson(chunkMessage.tool_calls),
                timestamp: Date.now(),
                images: [],
                tool_calls: chunkMessage.tool_calls, // 保存工具调用的方法和参数，方便后续使用，不用再向用户输出内容，只输出工具调用的方法和参数
            });
        } else if (lastMessage.role === ASSIST_ROLE) {
            lastMessage.content += chunkContent;
        } else {
            messages.push({
                id: createTimeStamp(),
                role: ASSIST_ROLE,
                content: chunkContent,
                timestamp: Date.now(),
                images: [],
            });
        }
        if (chunk.done) {
            state.isLoading = false;
        }
    });
    set((state) => {
        if (!chunk.done) {
            return;
        }
        const currentConversation = state.getCurrentConversation();
        if (!currentConversation) {
            throw new Error("no current conversation");
        }
        const messages = currentConversation.messages;
        const lastMessage = messages[messages.length - 1];
        if (lastMessage.role !== ASSIST_ROLE) {
            return;
        }
        if (lastMessage.content) {
            return;
        }
        const lastSecondMessage = messages[messages.length - 2];
        if (lastSecondMessage.role !== SYSTEM_ROLE) {
            return;
        }
        state.isLoading = true;
        sendMessage(
            (nextChunk: ChatResponse) => {
                onStream(nextChunk, set);
            },
            { messages: messages, tools: [] },
        );
    });
}

export const useChat = create(
    persist(chatStore, {
        name: "chat_history",
        partialize: (state) => ({
            currentConversationId: state.currentConversationId,
            conversations: state.conversations,
        }),
        storage: zustandPersistStorage,
        onRehydrateStorage: () => {
            return function (_, error) {
                useChat.setState({ isHydrated: true });
                if (!error) return;
                console.error("error: ", error);
            };
        },
    }),
);
